package io.hepu.robotize.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.aop.interceptor.SimpleAsyncUncaughtExceptionHandler;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;

import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * Scheduling pool for tasks
 */
@Slf4j
@Configuration
@EnableScheduling
public class ScheduleConfiguration implements SchedulingConfigurer, AsyncConfigurer {

    private static final int THREAD_NUMBER = Runtime.getRuntime().availableProcessors() * 8;

    @Bean
    public ThreadPoolTaskScheduler taskScheduler() {
        ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler() {
            private static final long serialVersionUID = -1L;

            @Override
            public void destroy() {
                /*
                 * Sets the policy on whether to execute existing delayed tasks even when this executor has been shutdown.
                 */
                this.getScheduledThreadPoolExecutor().setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
                super.destroy();
            }
        };
        scheduler.setPoolSize(THREAD_NUMBER);
        log.info("Scheduler pool with {} threads", THREAD_NUMBER);
        /*
         * Wait for all tasks to complete before continuing to destroy beans when thread pool is ready to close.
         */
        scheduler.setWaitForTasksToCompleteOnShutdown(true);
        /*
         * Set the maximum number of seconds for the task, if it exceeds this time, it will be forced to be destroyed,
         * to ensure thread pool can be shutdown in the end, rather than being blocked.
         */
        scheduler.setAwaitTerminationSeconds(30);
        scheduler.setThreadNamePrefix("Scheduler-");
        scheduler.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        return scheduler;
    }

    @Override
    public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
        TaskScheduler scheduler = this.taskScheduler();
        taskRegistrar.setTaskScheduler(scheduler);
    }

    @Override
    public Executor getAsyncExecutor() {
        return this.taskScheduler();
    }

    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return new SimpleAsyncUncaughtExceptionHandler();
    }

}
